# -*- coding: utf-8 -*-
# @Time    : 2021/7/7 14:46
# @Author  : LuoTianHang

# ####################datalist.py 说明##########################
# this script is used for the train.py to prepare the data

import os

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from config.config import BATCHSIZE
data_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
])


class StyleDataset(Dataset):
    def __init__(self, style_data_path="E:/coco_voc/train/"):
        super(StyleDataset, self).__init__()
        self.style_data_path = style_data_path
        self.data = self.getimages(self.style_data_path)

    def __getitem__(self, index):
        image = Image.open(self.style_data_path + self.data[index]).convert('RGB')
        image = data_transform(image)

        return image

    def __len__(self):
        return len(self.data)//BATCHSIZE*BATCHSIZE

    def getimages(self, datapath):
        data = []
        for image in os.listdir(datapath):
            data.append(image)
        return data
